import os
import h5py
import matplotlib.pyplot as plt
import numpy as np
from tqdm import tqdm
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
from torchvision.transforms import functional as F

def show_image(image, ind, epoch):
    image = image.permute(1, 2, 0)
    fig, axes = plt.subplots(1, 4, figsize=(16, 4))
    titles = ["Tracks", "ECAL", "HCAL"]

    for i in range(3):
        axes[i].imshow(image[:, :, i], cmap="viridis")
        axes[i].set_title(titles[i])
        axes[i].axis("off")

    axes[3].imshow(image * 255)
    axes[3].set_title("Combined")
    axes[3].axis("off")

    plt.tight_layout(pad=2.0)

    plt.show()
    plt.close(fig)

class CustomDataset(Dataset):
    def __init__(self, X, y):
        self.X = X
        self.y = y

        # Normalize each image so that the sum of pixel intensities is 1
        self.transform = transforms.Compose(
            [transforms.Lambda(lambda x: x / torch.sum(x, dim=(0, 1, 2)))]
        )

    def __len__(self):
        return len(self.X)

    def __getitem__(self, idx):
        image = self.X[idx]
        padded_image = np.zeros((128, 128, 3), dtype=np.float32)
        padded_image[: image.shape[0], : image.shape[1], :] = image

        return self.transform(
            torch.tensor(padded_image).permute(2, 0, 1)
        ), torch.tensor(self.y[idx])


class Autoencoder(nn.Module):
    def __init__(self):
        super(Autoencoder, self).__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(3, 16, 3, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(16, 32, 3, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(32, 64, 7),
        )
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(64, 32, 7),
            nn.ReLU(),
            nn.ConvTranspose2d(32, 16, 3, stride=2, padding=1, output_padding=1),
            nn.ReLU(),
            nn.ConvTranspose2d(16, 3, 3, stride=2, padding=1, output_padding=1),
            nn.Sigmoid(),
        )

    def forward(self, x):
        x = self.encoder(x)
        x = self.decoder(x)
        return x

with h5py.File(os.path.expanduser("~/data/quark-gluon_data-set_n139306.hdf5"), "r") as f:
    X = np.array(f["X_jets"])
    y = np.array(f["y"])

train_dataset = CustomDataset(X[:40000], y[:40000])
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

val_dataset = CustomDataset(X[40000:], y[40000:])
val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False)

model = Autoencoder()

criterion = nn.BCELoss()  # nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.0003)

random_indices = np.random.choice(len(val_dataset), size=20, replace=False)
random_images = torch.stack([val_dataset[i][0] for i in random_indices])

best_loss = 1
best_epoch = -1
num_epochs = 300
for epoch in range(num_epochs):
    print(f"epoch: {epoch+1}")
    running_loss = 0.0
    for images, _ in tqdm(train_loader):
        optimizer.zero_grad()
        reconstructions = model(images)
        loss = criterion(reconstructions, images)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    train_loss = running_loss / len(train_loader)
    print(f"train_loss: {train_loss}")

    """
    val_loss = 0.0
    for images, _ in val_loader:
        optimizer.zero_grad()
        reconstructions = model(images)
        loss = criterion(reconstructions, images) 
        val_loss += loss.item()
    print(f"val_loss: {running_loss / len(val_loader)}")
    """
    if not (epoch + 1) % 5:        
        with torch.no_grad():
            reconstructed_images = model(random_images)
            print('=' * 20)
            for i in range(5):
                print('Sample #%d' % i)
                show_image(random_images[i], i, -1)
                show_image(reconstructed_images[i], i, epoch)

    if train_loss < best_loss:
        best_loss = train_loss
        best_epoch = epoch
        torch.save(
            {"model_state_dict": model.state_dict()},
            "best_model.pth",
        )

    elif epoch - best_epoch > 10:
        print("Early stopped training at epoch %d" % epoch)
        break

epoch: 1

100%|█████████████████████████████████████████| 625/625 [28:39<00:00,  2.75s/it]

train_loss: 0.03246450614165515
epoch: 2

100%|█████████████████████████████████████████| 625/625 [40:52<00:00,  3.92s/it]

train_loss: 0.0019507485050708055
epoch: 3

100%|█████████████████████████████████████████| 625/625 [42:18<00:00,  4.06s/it]

train_loss: 0.0019496464380994438
epoch: 4

100%|█████████████████████████████████████████| 625/625 [38:43<00:00,  3.72s/it]

train_loss: 0.0019154714060947299
epoch: 5

100%|█████████████████████████████████████████| 625/625 [19:47<00:00,  1.90s/it]
Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers).